/*
 *
 * Copyright (C) 2008 Ed Hong
 * 
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */


#include "SubbandIterator.h"
#include "contexttree.h"

static void bubbledown(SubbandIterator*si, int ind);
static void updateHeap(SubbandIterator*si, int ind);

// standard order is
//  by bitplane
//   by subband
//    by component
//     by row 
//      by col
static void standard_order_advance(const SubbandIterator *si,  sbploc *sl) {
  sl->j++;
  if (sl->j == si->sc[sl->comp].sub[sl->subband].num_cols) {
    sl->j = 0;
    sl->i++;
    if (sl->i == si->sc[sl->comp].sub[sl->subband].num_rows) {
      sl->i = 0;
      sl->comp++;
      if (sl->comp == si->numComponents) {
	sl->comp = 0;
	sl->subband++;
	if (sl->subband == si->sc[sl->comp].numSubbands) {
	  sl->subband = 0;
	  sl->k++;
	}
      }
    }
  }
}

static void set_sbploc_zero(sbploc *s) {
  s->i = 0;
  s->j = 0;
  s->k = 0;
  s->subband = 0;
  s->comp = 0;
}



// updates context of coeff at sbploc
// returns 1 if context changed,
// returns 0 if it stays the same. 
static int update_context(CoeffHeapNode *chn, const SubbandComponent *sc, 
			  const contexttree *cc) {
  int newcontext;
  if (cc == NULL) {
    newcontext = 0;
  }
  else  {
        if (chn->loc.i == 1 && chn->loc.j == 6 && chn->loc.subband == 1 && chn->loc.k==4) {
	  //	  printf("hello2\n");
       }
    newcontext = find_context_num(sc, chn->loc.comp, chn->loc.i,
						  chn->loc.j, chn->loc.subband, 
						  chn->loc.k, cc);
  }
  // note that chn->loc.k must always == the coeff's currbitplane for this
  // to work. 

  if (chn->context == newcontext)
    return 0;
  
  chn->context = newcontext;
  return 1;

}

// resets all coeffss to start at bitplane 0.
static void reset_currbitplane(SubbandIterator *si) {
  //  printf("reset: %X \n", si);
    sbploc newloc;
    set_sbploc_zero(&newloc);
    
    while (newloc.k != 1) { 

      si->sc[newloc.comp].sub[newloc.subband].coeff[newloc.i][newloc.j].currbitplane = 0;

      standard_order_advance(si,&(newloc));
    }
}


// does two things, to update node at location (r,c). 
  //   (context update, and replacekey in heap ). 
// does not assume neighbor exists, but does assume 
// heap and sorting functions exist. 
static void updateNeighbor(SubbandIterator *si, int cn, int s, int r, int c) {
  SubbandBits *sb = &(si->sc[cn].sub[s]);
  if (coeff_exists(sb,r,c)) {
    if (update_context(si->dat[*(sb->coeff[r][c].locator)], si->sc, si->ct)) {
	    // replace_key check here. 
	    updateHeap(si, *(sb->coeff[r][c].locator));
    }
  }
}

	/* to update insignificance levels on the fly here....
	   using something like this. 
	  //recursivel check parents of this node 
	  int psub, pr, pc;
	  computeParent(s,r,c,&psub, &pr, &pc);
	  updateNeighbor(si, cn, psub, pr, pc, bp, neighborflag);	  
	  }*/


static void updateNeighbors8Adj( SubbandIterator *si, int cn, 
				      int r, int c, int s) {
  updateNeighbor(si, cn, s,r-1,c-1);
  updateNeighbor(si, cn, s,r-1,c);
  updateNeighbor(si, cn, s,r-1,c+1);
  updateNeighbor(si, cn, s,r,c-1);
  updateNeighbor(si, cn, s,r,c+1);
  updateNeighbor(si, cn, s,r+1,c-1);
  updateNeighbor(si, cn, s,r+1,c);
  updateNeighbor(si, cn, s,r+1,c+1);
}

static void updateNeighborsChildren(SubbandIterator *si, int cn,  
					 int r, int c, int s) {

  int childsub, childr, childc;
  computeChild(s,r,c,&childsub, &childr, &childc);    
  if (childsub < si->sc[cn].numSubbands && childsub > 0) {
    // child may exist. 
    updateNeighbor(si, cn, childsub,childr+1,childc+1);
    updateNeighbor(si, cn, childsub,childr+1,childc);
    updateNeighbor(si, cn, childsub,childr,childc+1);
    updateNeighbor(si, cn, childsub,childr,childc);
  }
}

static void updateNeighborsSubbandNeighbor(SubbandIterator *si, int cn,  
						int r, int c, int s) 
{
  // LL subband has no subband neighbors.
  if (s>0) {
    //spatially collocated neighbor at same subband level. 
    if (s%3 == 1) {
      updateNeighbor(si, cn, s+1,r,c);
      updateNeighbor(si, cn, s+2,r,c);
    }
    else if (s%3 ==2) {
      updateNeighbor(si, cn, s-1,r,c);
      updateNeighbor(si, cn, s+1,r,c);
    }
    else {
      updateNeighbor(si, cn, s-2,r,c);
      updateNeighbor(si, cn, s-1,r,c);
    }  
  }
}

static void updateNeighborsParent(SubbandIterator *si, int cn,  
				       int r, int c, int s)  {
  int psub, pr, pc;
  computeParent(s,r,c,&psub, &pr, &pc);
  updateNeighbor(si, cn, psub,pr,pc);
}

void updateNeighbors(SubbandIterator *si, int cn, 
			  int r, int c, int s){
  updateNeighbors8Adj(si,cn,r,c,s);
  updateNeighborsSubbandNeighbor(si,cn,r,c,s);
  updateNeighborsParent(si,cn,r,c,s);
  updateNeighborsChildren(si,cn,r,c,s);
}

static SubbandComponent *storedcomp;  // bad global-like variable. 
// stores most-recently used subbandcomp; sorting method can access it. 
// this means sorting will not work when multiple images are 
// simulatenously compressed. 


int SubbandIterator_InitAlloc(SubbandIterator *si, SubbandComponent *sc, 
				     int numComponents, 
				     int (*f)(const CoeffHeapNode *, 
					      const CoeffHeapNode *, 
					      const QccENTArithmeticModel *), 
			      contexttree *c, 
			      QccENTArithmeticModel *m
			      ) {
  //  printf("InitAlloc: %X %X %d %X %X\n", si, sc, numComponents, f, c);
  storedcomp = sc;
  si->dat = NULL;
  si->IsLessthan = f;
  set_sbploc_zero(&si->loc);
  si->numComponents = numComponents;
  si->sc = sc; 
  si->ct = c;
  si->totalcoeffs = 0;
  si->model = m;
  reset_currbitplane(si);
  if ((c==NULL && f !=NULL) ||
      (c!=NULL && f == NULL)) {
	QccErrorAddMessage("SubbandIterator_InitAlloc: Both function and context tree must be non-NULL for sorting");
	si->ct = NULL;
	si->IsLessthan = NULL;
	return 1;
  }
  if (c != NULL) { // initialize heap. 
    int totalcoeffs = sc[0].numCoeffs;
    int i;
    for (i=1; i < numComponents; i++) {
      totalcoeffs += sc[i].numCoeffs;
    }
    si->totalcoeffs = totalcoeffs;
    si->dat = (CoeffHeapNode **)malloc (sizeof(CoeffHeapNode *) * (totalcoeffs+1));
    
    if (si->dat == NULL) {
      QccErrorAddMessage("SubbandIterator_InitAlloc: Error allocating memory");
      return 1;
    }
    int index;
    index = 1;
    sbploc newloc;
    set_sbploc_zero(&newloc);
    
    while (newloc.k != 1) { // initialize heapnodes & locator pointers.
      si->dat[index] = (CoeffHeapNode *) malloc(sizeof(CoeffHeapNode));
      if (si->dat[index] == NULL) {
	QccErrorAddMessage("SubbandIterator_InitAlloc: Error allocating memory");
	return 1;
      }
      si->dat[index]->index = index;
      si->dat[index]->loc = newloc;
      update_context(si->dat[index], si->sc, si->ct);
      sc[newloc.comp].sub[newloc.subband].coeff[newloc.i][newloc.j].locator = 
	&(si->dat[index]->index);
      
      index++;
      standard_order_advance(si,&(newloc));
    }
    if (index != totalcoeffs+1) {
      QccErrorAddMessage("SubbandIterator_InitAlloc: totalCoeffs value of %d inconsistent with final index of %d", totalcoeffs, index);
	return 1;      
    }
    // heapify 
    for (i=totalcoeffs/2; i>=1; i--) {
      bubbledown(si, i);
    }
    // set new top. 
    si->loc = si->dat[1]->loc;
  }

  return 0;
}

sbploc SubbandIterator_curr(SubbandIterator *si) {
  return si->loc;
}

static int localdebug = 0;

void SubbandIterator_updateAdvance(SubbandIterator *si) {
  sbploc *s;
  s = &(si->loc); // for faster access.   
  //    printf("subit ");   printsbploc(s);
  if (s->subband ==8 && s->i == 8 && s->j == 26 && s->k == 5) {
    //            printf("hello\n");
  }

  // increment currbitplane... do this before neighbor update
  // so neighbor's updating their context will see 
  // new value at this node. 
  CoefficientData *cdat; // also for faster access;
  cdat =  &( si->sc[s->comp].sub[s->subband].coeff[s->i][s->j]);
  cdat->currbitplane ++;
  if (si->ct != NULL) {
    si->dat[*(cdat->locator)]->loc.k ++;
  
    // context update and replacekey of current node. 
      update_context(si->dat[*(cdat->locator)], si->sc, si->ct);
      // replace_key check here. 
    updateHeap(si, *(cdat->locator));
  

  // update neighbors; for each neighbor do..
  //   ( context update, and replacekey in heap )
    updateNeighbors(si, s->comp, s->i, s->j, s->subband);
    // no updating of neighbor contexts necessary if we are not using sorting. 
    //   even though some characteristics computed on the fly
    //   contextnums are always completely recomputed with current 
    //   info when figuring out context of next bit right before encode/decode. 
  }
 
  // advance to next item
  if (si->IsLessthan == NULL) {
    standard_order_advance(si, s);
    //    printsbploc(s);
  }
  else { // find next root node of heap. 
    // re-heapify: should only make a difference when 
    //  using ordering that depends on frequency of 0's in a context; many
    //  coeffs with the same context as context used in coding of last bit
    //    could potentially have to change spots. 
  /*  int i; // reheapifying too expensive computationally to run after every bit. 
    localdebug = 1;
    for (i=si->totalcoeffs/2; i>=1; i--) {
      bubbledown(si, i);
    }
    localdebug = 0;
*/   
    // set next
    si->loc = si->dat[1]->loc;
  }

}

void SubbandIterator_free(SubbandIterator *si) {
  if (si->ct != NULL) {
    if (si->dat != NULL) {
      int index;
      for(index = 1; index < si->sc->numCoeffs+1; index++) {
	if (si->dat[index] != NULL) {
	  free(si->dat[index]);
	}
      }
      free(si->dat);
      si->dat = NULL;
    }
    si->ct = NULL;
  }
  else if (si->dat != NULL) {
    printf("SubbandIterator_free: error; dat is not NULL but should be\n");
  }

} 


static void swap(SubbandIterator*si, int index1, int index2) {
  if (localdebug==1) printf("^");
  //  if(si->loc.k==5 && (index1==76329) && index2 ==38164)
  //  printf("subitswap: %d %d\n", index1, index2);
  CoeffHeapNode *temp = si->dat[index1];
  si->dat[index1] = si->dat[index2];
  si->dat[index2] = temp;
  si->dat[index1]->index = index1;
  si->dat[index2]->index = index2;
}

static void bubbleup(SubbandIterator*si, int ind) {
  if (ind<=1) return;

  if ((*(si->IsLessthan))(si->dat[ind], si->dat[ind/2], si->model)) {
    swap(si, ind, ind/2);
    bubbleup(si, ind/2);
  }
}

static void bubbledown(SubbandIterator*si, int ind) {
  if (2*ind > si->totalcoeffs) {
    return;
  }
  else if (2*ind == si->totalcoeffs) {
    // only one child, do comparison & possible swap. 
    if ((*(si->IsLessthan))(si->dat[2*ind], si->dat[ind], si->model)) {
      swap(si, ind, 2*ind);
    }
  }
  else {
    int smallerchildindex; 
    if ((*(si->IsLessthan))(si->dat[2*ind+1], si->dat[2*ind], si->model)) {
      smallerchildindex = 2*ind+1;
    }
    else smallerchildindex = 2*ind;
    if ((*(si->IsLessthan))(si->dat[smallerchildindex], si->dat[ind], si->model)) {
      swap(si, smallerchildindex, ind);
      bubbledown(si, smallerchildindex);
    }    
  }  
}

static void updateHeap(SubbandIterator*si, int ind) {
  if (ind >1) {
    // there is a parent, check for bubbleup. 
    if ((*(si->IsLessthan))(si->dat[ind], si->dat[ind/2], si->model)) {
      bubbleup(si, ind);
      return;
    }    
  }
  bubbledown(si,ind);
}


int standard_order_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  }
  else if (x1->loc.k == x2->loc.k){
    if (x1->loc.subband < x2->loc.subband) {
      return 1; 
    }
    else  if (x1->loc.subband == x2->loc.subband) {
      if (x1->loc.comp < x2->loc.comp) {
	return 1; 
      }
      else  if (x1->loc.comp == x2->loc.comp) {
	if (x1->loc.i < x2->loc.i) {
	  return 1; 
	}
	else  if (x1->loc.i == x2->loc.i) {
	  if (x1->loc.j < x2->loc.j) {
	    return 1; 
	  }
	}	
      }
    }
  }
  
  return 0;
}

int reverse_standard_order_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else return standard_order_less_than(x2,x1, m);
}

int by_context_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  return x1->context < x2->context;

}


int reverse_by_context_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  return by_context_less_than(x2,x1, m);
}

int reverse_by_context_then_standard_order_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else if (x1->context == x2->context) {
    return standard_order_less_than(x1,x2, m);
  }
  else
    return by_context_less_than(x2,x1, m);
}

// compares frequncy of context; returns -1, 0, or 1 for <, ==, >  of x1ratio op x2ratio
// ratios are computed as  frequency_ones / frequency_zeros.  (more ones => higher ratio)
static int freq_of_context_compare(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
    int x1freq0, x1freq1, x2freq0, x2freq1;

    if (x1->context == x2->context) return 0;

    x1freq0 = QccENTArithmeticGetSymbolFrequency(m, x1->context, 0);
    x1freq1 = QccENTArithmeticGetSymbolFrequency(m, x1->context, 1);
    x2freq0 = QccENTArithmeticGetSymbolFrequency(m, x2->context, 0);
    x2freq1 = QccENTArithmeticGetSymbolFrequency(m, x2->context, 1);

    /* actually, we want to always encode 1's before zeroes, 
       even if there are many 1's in a row, so ignore this part. 
    // use freq0 for more probable symbol, freq1 for less probable symbol.
    if (x1freq0 < x1freq1) {
      int temp;
      temp = x1freq0;
      x1freq0 = x1freq1;
      x1freq1 = temp;
    }
    if (x2freq0 < x2freq1) {
      int temp;
      temp = x2freq0;
      x2freq0 = x2freq1;
      x2freq1 = temp;
    }
    */
    double x1ratio, x2ratio;
    //    frequencies guaranteed to be nonzero.
    x1ratio = (double) x1freq1 / (double) x1freq0; 
    x2ratio = (double) x2freq1 / (double) x2freq0;
    // ratios should be between 0 and freq1, and
    // rarely bigger than 1 (we are compressing sequences of many 0's). 
    if (x1ratio <  x2ratio) return -1;
    else if (x1ratio > x2ratio) return 1;
    else return 0;
}

int by_freq_of_context_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m)  {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else {
    return (freq_of_context_compare(x1,x2,m)==1);
  }
}

int by_freq_of_context_then_standard_order_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m)  {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else {
    int freqcompare = freq_of_context_compare(x1,x2,m);
    if (freqcompare==0)
      return standard_order_less_than(x1,x2, m);
    else return (freqcompare == 1);
  }

}



int by_issig_then_freq_of_context_then_standard_order_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m)  {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else {
    int x1sig, x2sig;
    x1sig = is_sig(&(storedcomp[x1->loc.comp].sub[x1->loc.subband]),
		   x1->loc.i, x1->loc.j, x1->loc.k);
    x2sig = is_sig(&(storedcomp[x2->loc.comp].sub[x2->loc.subband]),
		   x2->loc.i, x2->loc.j, x2->loc.k);
    if (x1sig > x2sig)
      return 1;
    else if (x2sig > x1sig)
      return 0;
    else return by_freq_of_context_then_standard_order_less_than(x1, x2, m) ;

  }
}

// returns -1, 0, or 1 for <, ==, >  of x1 compare x2; lessthan means should appear first in bitstream.
int JPEG2k_3_fractional_bitplanes_compare(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  if (x1->loc.k < x2->loc.k) { // first by bitplane. 
    return -1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 1;
  } 
  else { // now in order significance bits w/ >=1 sig neighbor, 
    // then refinement
    // then rest of significance bits. 
    int x1sig, x2sig;
    x1sig = is_sig(&(storedcomp[x1->loc.comp].sub[x1->loc.subband]),
		   x1->loc.i, x1->loc.j, x1->loc.k);
    x2sig = is_sig(&(storedcomp[x2->loc.comp].sub[x2->loc.subband]),
		   x2->loc.i, x2->loc.j, x2->loc.k);
    if (!x1sig && !x2sig)
        return 0; // refinement bits are equivalent. 
    
    int x1sigcount, x2sigcount;
    if (!x1sig) {
       x2sigcount = sig_count_8Neighb(&(storedcomp[x2->loc.comp]),
		   x2->loc.i, x2->loc.j, x2->loc.subband, x2->loc.k);
       if (x2sigcount >0) return 1;
       else return -1;
    }
    else if (!x2sig) {
       x1sigcount = sig_count_8Neighb(&(storedcomp[x1->loc.comp]),
		   x1->loc.i, x1->loc.j, x1->loc.subband, x1->loc.k);
       if (x1sigcount >0) return -1;
       else return 1;
    
    }
    else {
       x1sigcount = sig_count_8Neighb(&(storedcomp[x1->loc.comp]),
		   x1->loc.i, x1->loc.j, x1->loc.subband, x1->loc.k);
       x2sigcount = sig_count_8Neighb(&(storedcomp[x2->loc.comp]),
		   x2->loc.i, x2->loc.j, x2->loc.subband, x2->loc.k);
       return (x1sigcount ==0) - (x2sigcount==0);
    }  
  }
}


int JPEG2k_3_fractional_bitplanes_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
   return (JPEG2k_3_fractional_bitplanes_compare(x1,x2,m) == -1);
}

int by_JPEG2k_3_fractional_bitplanes_then_freq_and_standard_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m) {
  int fractionalcompare = JPEG2k_3_fractional_bitplanes_compare(x1, x2, m); 
  if (fractionalcompare <0) 
    return 1; 
  else if (fractionalcompare > 0)
    return 0;
  else {
     return by_freq_of_context_then_standard_order_less_than(x1, x2, m) ;
  }
}

int (*get_sorting_function(int num))(const CoeffHeapNode *, const CoeffHeapNode *, const QccENTArithmeticModel *) {
  switch (num) {
  case 0: return NULL;
  case 1: return standard_order_less_than;
  case 2: return reverse_standard_order_less_than;
  case 3: return by_context_less_than;
  case 4: return reverse_by_context_less_than;
  case 5: return reverse_by_context_then_standard_order_less_than;
  case 6: return by_freq_of_context_less_than;
  case 7: return by_freq_of_context_then_standard_order_less_than;
  case 8: return by_issig_then_freq_of_context_then_standard_order_less_than;
  case 9: return JPEG2k_3_fractional_bitplanes_less_than;
  case 10: return by_JPEG2k_3_fractional_bitplanes_then_freq_and_standard_less_than;
  default: 
    printf("get_sorting_function: unrecognized sorting function type... exiting\n");
    exit(1);
  }
}

int by_freq_of_context_then_pattern_less_than(const CoeffHeapNode *x1, const CoeffHeapNode*x2, const QccENTArithmeticModel *m)  {
  if (x1->loc.k < x2->loc.k) {
    return 1;
  } 
  else if (x1->loc.k > x2->loc.k) {
    return 0;
  } 
  else if (x1->context == x2->context) {
    return standard_order_less_than(x1,x2, m);
  }
  else
    return by_freq_of_context_less_than(x1,x2, m);
}
